import face_recognition
import cv2
import numpy as np
import os
from typing import List, Tuple
from src.utils.face_utils import put_chinese_text


class FaceRecognition:
    """
    人脸识别类，提供人脸检测、编码和识别功能
    """

    def __init__(self):
        """
        初始化人脸识别器
        """
        self.known_face_encodings = []
        self.known_face_names = []

    def load_known_faces(self, faces_dir: str) -> bool:
        """
        从指定目录加载已知人脸图像

        Args:
            faces_dir (str): 包含已知人脸图像的目录路径

        Returns:
            bool: 是否成功加载至少一张人脸
        """
        if not os.path.exists(faces_dir):
            print(f"警告: 目录 {faces_dir} 不存在")
            return False

        loaded_count = 0
        for filename in os.listdir(faces_dir):
            if filename.lower().endswith(('.jpg', '.png', '.jpeg')):
                image_path = os.path.join(faces_dir, filename)
                try:
                    image = face_recognition.load_image_file(image_path)
                    face_encodings = face_recognition.face_encodings(image)

                    if len(face_encodings) > 0:
                        self.known_face_encodings.append(face_encodings[0])
                        # 使用文件名作为人名（去掉扩展名）
                        name = os.path.splitext(filename)[0]
                        self.known_face_names.append(name)
                        loaded_count += 1
                        print(f"成功加载人脸: {name}")
                    else:
                        print(f"警告: 在 {filename} 中未检测到人脸")
                except Exception as e:
                    print(f"加载 {filename} 时出错: {e}")

        print(f"总共加载了 {loaded_count} 张人脸")
        return loaded_count > 0

    def recognize_from_image(self, image_path: str) -> Tuple[List[Tuple[int, int, int, int]], List[str]]:
        """
        从图片中识别人脸

        Args:
            image_path (str): 待识别图片的路径

        Returns:
            tuple: (人脸位置列表, 对应的人名列表)
        """
        if not os.path.exists(image_path):
            raise FileNotFoundError(f"图片文件 {image_path} 不存在")

        # 加载图片
        image = face_recognition.load_image_file(image_path)
        rgb_image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)

        # 找到图片中所有人脸的位置和编码
        face_locations = face_recognition.face_locations(rgb_image)
        face_encodings = face_recognition.face_encodings(rgb_image, face_locations)

        face_names = []
        for face_encoding in face_encodings:
            # 与已知人脸对比
            matches = face_recognition.compare_faces(self.known_face_encodings, face_encoding)
            name = "Unknown"

            # 计算距离，找到最匹配的人脸
            if len(self.known_face_encodings) > 0:
                face_distances = face_recognition.face_distance(self.known_face_encodings, face_encoding)
                best_match_index = np.argmin(face_distances)

                if matches[best_match_index]:
                    name = self.known_face_names[best_match_index]

            face_names.append(name)

        return face_locations, face_names

    def recognize_from_video(self, video_source=0) -> None:
        """
        从摄像头或视频文件实时识别人脸

        Args:
            video_source: 视频源，可以是摄像头索引(如0)或视频文件路径
        """
        video_capture = cv2.VideoCapture(video_source)

        print("开始人脸识别，按 'q' 键退出")
        while True:
            ret, frame = video_capture.read()
            if not ret:
                print("无法读取视频帧")
                break

            # 缩小帧尺寸以提高处理速度
            small_frame = cv2.resize(frame, (0, 0), fx=0.25, fy=0.25)
            rgb_small_frame = cv2.cvtColor(small_frame, cv2.COLOR_BGR2RGB)

            # 找到当前帧中所有人脸位置和编码
            face_locations = face_recognition.face_locations(rgb_small_frame)
            face_encodings = face_recognition.face_encodings(rgb_small_frame, face_locations)

            face_names = []
            for face_encoding in face_encodings:
                matches = face_recognition.compare_faces(self.known_face_encodings, face_encoding)
                name = "Unknown"

                if len(self.known_face_encodings) > 0:
                    face_distances = face_recognition.face_distance(self.known_face_encodings, face_encoding)
                    best_match_index = np.argmin(face_distances)

                    if matches[best_match_index]:
                        name = self.known_face_names[best_match_index]
                face_names.append(name)

            # 显示结果
            for (top, right, bottom, left), name in zip(face_locations, face_names):
                # 放大坐标回到原始帧尺寸
                top *= 4
                right *= 4
                bottom *= 4
                left *= 4

                # 绘制人脸框
                cv2.rectangle(frame, (left, top), (right, bottom), (0, 0, 255), 2)
                cv2.rectangle(frame, (left, bottom - 35), (right, bottom), (0, 0, 255), cv2.FILLED)
                
                # 使用PIL处理中文显示
                frame = put_chinese_text(frame, name, (left + 6, bottom - 30), 20, (255, 255, 255))

            cv2.imshow('Face Recognition', frame)

            # 按q键退出
            key = cv2.waitKey(1) & 0xFF
            if key == ord('q') or key == ord('Q') or key == 27:  # 27是ESC键的ASCII码
                break

        video_capture.release()
        cv2.destroyAllWindows()

    def add_face(self, image_path: str, name: str) -> bool:
        """
        添加新的人脸到已知人脸库

        Args:
            image_path (str): 人脸图片路径
            name (str): 人名

        Returns:
            bool: 是否成功添加
        """
        if not os.path.exists(image_path):
            print(f"错误: 图片文件 {image_path} 不存在")
            return False

        try:
            # 使用numpy处理可能包含中文的路径
            with open(image_path, 'rb') as f:
                file_bytes = np.asarray(bytearray(f.read()), dtype=np.uint8)
                image = cv2.imdecode(file_bytes, cv2.IMREAD_COLOR)
            if image is None:
                print(f"错误: 无法读取图像文件 {image_path}")
                return False
                
            # 转换BGR到RGB
            rgb_image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
            face_encodings = face_recognition.face_encodings(rgb_image)

            if len(face_encodings) > 0:
                self.known_face_encodings.append(face_encodings[0])
                self.known_face_names.append(name)
                print(f"成功添加人脸: {name}")
                return True
            else:
                print(f"在 {image_path} 中未检测到人脸")
                return False
        except Exception as e:
            print(f"添加人脸时出错: {e}")
            return False
